# ===----------------------------------------------------------------------=== #
# Copyright (c) 2025, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #

from pathlib import Path
from sys._assembly import inlined_assembly

from gpu import barrier, thread_idx
from gpu.host import DeviceContext
from gpu.host.compile import _compile_code
from memory import LegacyUnsafePointer as UnsafePointer, stack_allocation


fn kernel(x: Int) -> Int:
    return Int(thread_idx.x)


fn parametric[f: fn (Int) -> Int]() -> Int:
    return f(42)


# CHECK-LABEL: test_compile_code
def test_compile_code():
    print("== test_compile_code")

    # CHECK: Generated by LLVM NVPTX Back-End
    # CHECK: .visible .entry  test_compile_nvptx_kernel_
    print(_compile_code[kernel]())
    # CHECK: Generated by LLVM NVPTX Back-End
    # CHECK: .visible .entry test_compile_nvptx_parametric_
    print(_compile_code[parametric[kernel]]())
    # CHECK: ; ModuleID =
    # CHECK: lvm.nvvm.read.ptx.sreg.tid.x()
    print(_compile_code[parametric[kernel], emission_kind="llvm"]())


# CHECK-LABEL: test_compile_function
def test_compile_function():
    print("== test_compile_function")

    fn kernel(x: UnsafePointer[Int]):
        x[0] = Int(thread_idx.x)

    # CHECK: tid.x

    with DeviceContext() as ctx:
        _ = ctx.compile_function[kernel, dump_asm=True]()


fn kernel_inlined_assembly():
    inlined_assembly["nanosleep.u32 $0;", NoneType, constraints="r"](
        UInt32(100)
    )


# CHECK-LABEL: test_compile_function_with_assembly
def test_compile_function_with_assembly():
    print("== test_compile_function_with_assembly")

    # CHECK: nanosleep
    # CHECK-NOT: begin assembly

    with DeviceContext() as ctx:
        _ = ctx.compile_function[kernel_inlined_assembly, dump_asm=True]()


# CHECK-LABEL: test_compile_function_with_path
def test_compile_function_with_path():
    print("== test_compile_function_with_path")

    # CHECK: nanosleep
    # CHECK-NOT: begin assembly

    with DeviceContext() as ctx:
        alias out_file = Path("/tmp/my_file.ptx")
        _ = ctx.compile_function[kernel_inlined_assembly, dump_asm=out_file]()
        print(out_file.read_text())


# CHECK-LABEL: test_compile_function_with_path_func
def test_compile_function_with_path_func():
    print("== test_compile_function_with_path_func")

    # CHECK: nanosleep
    # CHECK-NOT: begin assembly

    with DeviceContext() as ctx:
        var out_file_name = "my_file_2.ptx"
        alias out_dir = Path("/tmp")

        @parameter
        fn dummy_fn() capturing -> Path:
            return out_dir / out_file_name

        _ = ctx.compile_function[kernel_inlined_assembly, dump_asm=dummy_fn]()

        var out_file = out_dir / out_file_name
        print(out_file.read_text())


# CHECK-LABEL: test_short_nvptx_ptr
def test_short_nvptx_ptr():
    print("== test_short_nvptx_ptr")

    fn do_some_shared_mem_op(src: UnsafePointer[Int32]):
        var a = stack_allocation[
            20, Int32, address_space = AddressSpace.SHARED
        ]()
        a[thread_idx.x] = src[0]
        barrier()

    # CHECK: // demoted variable
    # CHECK-NOT: mov.u64
    # CHECK: mov.u32
    # CHECK-NEXT: shl.b32
    # CHECK-NEXT: mov.b32
    # CHECK-NEXT: add.s32
    # CHECK-NEXT: ld.global.b32
    # CHECK-NEXT: st.shared.b32
    with DeviceContext() as ctx:
        _ = ctx.compile_function[do_some_shared_mem_op, dump_asm=True]()


def main():
    test_compile_code()
    test_compile_function()
    test_compile_function_with_assembly()
    test_compile_function_with_path()
    test_compile_function_with_path_func()
    test_short_nvptx_ptr()
